package com.one.blocks.jdbc.meta.dialect;

import com.one.blocks.jdbc.meta.CatalogMeta;
import com.one.blocks.jdbc.meta.ColumnMeta;
import com.one.blocks.jdbc.meta.TableMeta;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.function.Predicate;

/**
 * @author <a href="mailto:idler41@163.con">linfuxin</a> created on 2024-07-16 20:22:08
 */
public class PostgresqlDialect implements MetaDialect {
    private static final String COLUMN_SQL = "SELECT\n" +
            "    attname AS COLUMN_NAME,\n" +
            "    format_type(atttypid, atttypmod) AS COLUMN_TYPE\n" +
            "FROM\n" +
            "    pg_attribute\n" +
            "JOIN\n" +
            "    pg_class ON pg_attribute.attrelid = pg_class.oid\n" +
            "JOIN\n" +
            "    pg_namespace ON pg_class.relnamespace = pg_namespace.oid\n" +
            "JOIN\n" +
            "    pg_type ON pg_attribute.atttypid = pg_type.oid\n" +
            "WHERE\n" +
            "    pg_namespace.nspname = ?\n" +
            "    AND pg_class.relname = ?\n" +
            "    AND attnum > 0 AND NOT attisdropped\n" +
            "ORDER BY attnum";

    private static final Set<String> SYSTEM_SCHEMA = new HashSet<>(Arrays.asList("information_schema", "pg_catalog"));

    private static final String SUPPORT = "postgresql";

    @Override
    public String support() {
        return SUPPORT;
    }

    @Override
    public List<CatalogMeta> getCatalog(DataSource ds) {
        List<CatalogMeta> data = new ArrayList<>();
        try (Connection con = ds.getConnection()) {
            DatabaseMetaData metaData = con.getMetaData();
            try (ResultSet resultSet = metaData.getSchemas()) {
                while (resultSet.next()) {
                    // 过滤系统schema
                    String schema = resultSet.getString("TABLE_SCHEM");
                    if (SYSTEM_SCHEMA.contains(schema)) {
                        continue;
                    }
                    CatalogMeta entity = new CatalogMeta();
                    entity.setCatalogName(schema);
                    data.add(entity);
                }
            }
        } catch (SQLException e) {
            throw new RuntimeException(e);
        }
        return data;
    }

    @Override
    public List<TableMeta> getTable(DataSource ds, String catalog, Predicate<String> filter) {
        List<TableMeta> data = new ArrayList<>();
        try (Connection con = ds.getConnection()) {
            DatabaseMetaData metaData = con.getMetaData();
            // 传入的catalog其实是schema
            try (ResultSet rs = metaData.getTables(null, catalog, "%", new String[]{"TABLE"})) {
                while (rs.next()) {
                    String tableName = rs.getString("TABLE_NAME");
                    if (filter == null || filter.test(tableName)) {
                        TableMeta item = new TableMeta();
                        item.setTableName(tableName);
                        item.setTableComment(rs.getString("REMARKS"));
                        data.add(item);
                    }
                }
            }
        } catch (SQLException e) {
            throw new RuntimeException(e);
        }
        return data;
    }

    @Override
    public List<ColumnMeta> getColumn(DataSource ds, String catalog, String tableName, boolean needColumnType) {
        Map<String, ColumnMeta> data = new LinkedHashMap<>();
        try (Connection con = ds.getConnection()) {
            DatabaseMetaData metaData = con.getMetaData();
            try (ResultSet columnResult = metaData.getColumns(catalog, null, tableName, null);
                 ResultSet pkResultSet = metaData.getPrimaryKeys(catalog, null, tableName)) {
                Set<String> pkColumnName = new LinkedHashSet<>();
                while (pkResultSet.next()) {
                    pkColumnName.add(pkResultSet.getString(4));
                }
                while (columnResult.next()) {
                    ColumnMeta columnMeta = new ColumnMeta();
                    columnMeta.setDbName(catalog);
                    columnMeta.setTableName(tableName);
                    columnMeta.setColumnName(columnResult.getString("COLUMN_NAME"));
                    columnMeta.setTypeName(columnResult.getString("TYPE_NAME"));
                    columnMeta.setComment(columnResult.getString("REMARKS"));
                    columnMeta.setJdbcType(columnResult.getInt("DATA_TYPE"));
                    columnMeta.setPrecision(columnResult.getInt("COLUMN_SIZE"));
                    columnMeta.setScale(columnResult.getInt("DECIMAL_DIGITS"));
                    // jdbc规范IS_NULLABLE列返回Int,但很多主流数据库并没有遵守规范，比如mysql、pgsql、oracle返回字符串"NO"(不同版本可能有差异)
                    columnMeta.setNullable(columnResult.getObject("IS_NULLABLE"));
                    columnMeta.setIdentity(pkColumnName.contains(columnMeta.getColumnName()));
                    data.put(columnMeta.getColumnName(), columnMeta);
                }

                if (needColumnType) {
                    try (PreparedStatement stmt = con.prepareStatement(COLUMN_SQL)) {
                        // TODO 根据catalog查询所有schema
                        stmt.setString(1, catalog);
                        stmt.setString(2, tableName);
                        try (ResultSet colTypeResult = stmt.executeQuery()) {
                            while (colTypeResult.next()) {
                                ColumnMeta columnMeta = data.get(colTypeResult.getString("COLUMN_NAME"));
                                if (columnMeta != null) {
                                    columnMeta.setColumnType(colTypeResult.getString("COLUMN_TYPE"));
                                }
                            }
                        }
                    }
                }
            }
        } catch (SQLException e) {
            throw new RuntimeException(e);
        }
        return new ArrayList<>(data.values());
    }
}
